import math
import numpy as np
from core.cuda import cuda, cuda_module
from core.data.dataloader import DataLoader


class SeqDataLoader(DataLoader):
    """
    顺序数据加载器类，继承自 DataLoader。

    Args:
        dataset: 数据集，可迭代对象。
        batch_size (int): 批处理大小。
        gpu (bool, optional): 是否使用 GPU。默认为 False。
    """

    def __init__(self, dataset, batch_size, gpu=False):
        """
        初始化顺序数据加载器。

        Args:
            dataset: 数据集，可迭代对象。
            batch_size (int): 批处理大小。
            gpu (bool, optional): 是否使用 GPU。默认为 False。
        """
        super().__init__(dataset=dataset, batch_size=batch_size, shuffle=False, gpu=gpu)

    def __next__(self):
        """
        获取下一个批次的数据。

        Raises:
            StopIteration: 当所有批次都被处理后抛出。

        Returns:
            tuple: 包含输入数据和对应标签的元组。
        """
        if self.iteration >= self.max_iter:
            self.reset()
            raise StopIteration

        jump = self.data_size // self.batch_size
        # 计算当前批次的索引
        batch_index = [(i * jump + self.iteration) % self.data_size for i in range(self.batch_size)]
        # 获取当前批次的数据
        batch = [self.dataset[i] for i in batch_index]

        xp = cuda.cupy if self.gpu else np
        # 转换为数组
        x = xp.array([example[0] for example in batch])
        t = xp.array([example[1] for example in batch])

        self.iteration += 1
        return x, t
